{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "305ea615-b3aa-497b-866d-7536f3a4bfcb",
   "metadata": {},
   "source": [
    "# Headless Playground"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fd4006f-aa95-40fd-bcbb-6be2668d1a4d",
   "metadata": {},
   "source": [
    "*This mode is suitable to display the gui on remote machines (i.e. cannot set up the polyscope gui).*\n",
    "Here we show how to use `Engine3DGRUT` with [NVIDIA kaolin's[(https://github.com/NVIDIAGameWorks/kaolin) viewer. \n",
    "\n",
    "Note: Other viewers can similarly interact with the engine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb540372-a221-48f7-b8e9-c4d742d8ae7a",
   "metadata": {},
   "source": [
    "## Engine Params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfc2ea73-d020-4295-8bdc-9f30617c813b",
   "metadata": {},
   "source": [
    "Point at your playground params here.\n",
    "\n",
    "* *gs_object*: Path of pretrained 3dgrt checkpoint, as .pt / .ingp / .ply file.\n",
    "* *mesh_assets*: Path to folder containing mesh assets of .obj or .glb format.\n",
    "* *default_config*: Name of default config to use for .ingp, .ply files, or .pt files not trained with 3dgrt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf0bead-1fc0-421f-bf32-2da2d99d6e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_object = \"3dgrut/runs/<YOUR_MODEL_HERE>/ckpt_last.pt\"\n",
    "mesh_assets_folder = \"./assets\"\n",
    "default_config = \"apps/colmap_3dgrt.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efd77d91-94f7-4960-8388-068f3a94be07",
   "metadata": {},
   "source": [
    "## Install additional requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e9bbef1-d64e-4719-b89f-4cd562eb361b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install matplotlib ipywidgets --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1d345c3-5c58-4fcf-99e1-17b30c79515a",
   "metadata": {},
   "source": [
    "#### Add 3dgrut root to search path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "411aae24-bfac-4311-9949-d1da02bddf56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "notebook_path = Path().resolve()\n",
    "root_path = notebook_path.parent\n",
    "sys.path.append(str(root_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c02229-c167-417a-9343-efd8b45643d0",
   "metadata": {},
   "source": [
    "## 3DGRUT Playground Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31bd67bf-2348-431f-9ffe-34e5ea2b472b",
   "metadata": {},
   "source": [
    "#### Setup Headless Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a64151e22187546",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import copy\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.transforms.functional as F\n",
    "import kaolin\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from threedgrut.utils.logger import logger\n",
    "from threedgrut.gui.ps_extension import initialize_cugl_interop\n",
    "from threedgrut_playground.utils.video_out import VideoRecorder\n",
    "from threedgrut_playground.engine import Engine3DGRUT, OptixPrimitiveTypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e3640e5-0b1b-47de-b45d-e739a2edc5e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "engine = Engine3DGRUT(\n",
    "    gs_object=gs_object,\n",
    "    mesh_assets_folder=mesh_assets_folder,\n",
    "    default_config=default_config\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8057aa9-c118-4487-928a-e759e40e1979",
   "metadata": {},
   "source": [
    "#### Configure 3DGRUT Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ff54b67-615c-49aa-9448-7d823ca801bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure rendering settings\n",
    "engine.camera_type = 'Pinhole'\n",
    "engine.camera_fov = 60.0\n",
    "engine.use_spp = True\n",
    "engine.antialiasing_mode = '8x MSAA'\n",
    "\n",
    "# Remove initial glass sphere from scene\n",
    "for mesh_name in list(engine.primitives.objects.keys()):\n",
    "    engine.primitives.remove_primitive(mesh_name)\n",
    "\n",
    "# Add a glass 'Armadillo' to the scene\n",
    "# Get the asset with download_assets.sh script\n",
    "engine.primitives.add_primitive(\n",
    "    geometry_type='Armadillo',\n",
    "    primitive_type=OptixPrimitiveTypes.GLASS,\n",
    "    device='cuda'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6eab7bb-ed87-4319-8356-dae3f53a0a66",
   "metadata": {},
   "source": [
    "## Render Single Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea5c90be-9317-451c-9275-afcfe8dfda3f",
   "metadata": {},
   "source": [
    "The following block shows how to generate a single screenshot using the engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbc0c303-cd6a-41d9-9427-97395484212c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can also create a camera from (eye, at, up) or 4x4 view-matrix\n",
    "# See here: https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/recipes/camera/camera_init_simple.py\n",
    "\n",
    "# Create a camera programatically and position it\n",
    "camera = kaolin.render.easy_render.default_camera(512).cuda()\n",
    "camera.move_forward(-1.5)\n",
    "camera.move_up(1.5)\n",
    "camera.rotate(yaw=1.0, pitch=0.5, roll=2.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f3f71bf-f2b5-494b-a6e7-61b406059c03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Render a full quality frame\n",
    "framebuffer = engine.render(camera)\n",
    "rgba_buffer = torch.cat([framebuffer['rgb'], framebuffer['opacity']], dim=-1)\n",
    "\n",
    "# Display\n",
    "chw_buffer = rgba_buffer[0].permute(2, 0, 1)\n",
    "img = F.to_pil_image(chw_buffer)\n",
    "plt.imshow(img)\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43fd7f65-da85-4c2c-bd74-3267b7f6a845",
   "metadata": {},
   "source": [
    "## Interactive Renderer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c787a08a-8d90-4df5-ae8e-490d99184cc6",
   "metadata": {},
   "source": [
    "The following block shows how to use kaolin's internal viewer to drive the 3DGRUT engine with user interaction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53aa5e9a-cadb-489f-addc-3e735cb9a704",
   "metadata": {},
   "source": [
    "#### Set up widgets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50c550b-3726-4ab9-a628-24f7f50dcfd1",
   "metadata": {},
   "source": [
    "Set up some checkboxes, dropdowns and slider to control simple engine functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e3d0193-8dd0-4255-a294-a72025720473",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "\n",
    "aa_checkbox = widgets.Checkbox(\n",
    "    value=engine.use_spp,\n",
    "    description='Toggle Antialiasing'\n",
    ")\n",
    "\n",
    "aa_mode_combo = widgets.Dropdown(\n",
    "    options=engine.ANTIALIASING_MODES,\n",
    "    value='4x MSAA',\n",
    "    description='AA Mode'\n",
    ")\n",
    "\n",
    "denoiser_checkbox = widgets.Checkbox(\n",
    "    value=engine.use_optix_denoiser,\n",
    "    description='Toggle Optix Denoiser'\n",
    ")\n",
    "\n",
    "spp_slider = widgets.IntSlider(\n",
    "    value=engine.spp.spp,\n",
    "    min=1,\n",
    "    max=64,\n",
    "    step=1,\n",
    "    orientation='horizontal',\n",
    "    description='AA SPP',\n",
    "    disabled=(engine.spp.mode == 'msaa')\n",
    ")\n",
    "\n",
    "def on_change(change):\n",
    "    engine.use_spp = aa_checkbox.value\n",
    "    \n",
    "    engine.antialiasing_mode = aa_mode_combo.value\n",
    "    if engine.antialiasing_mode == '4x MSAA':\n",
    "        engine.spp.mode = 'msaa'\n",
    "        engine.spp.spp = 4\n",
    "    elif engine.antialiasing_mode == '8x MSAA':\n",
    "        engine.spp.mode = 'msaa'\n",
    "        engine.spp.spp = 8\n",
    "    elif engine.antialiasing_mode == '16x MSAA':\n",
    "        engine.spp.mode = 'msaa'\n",
    "        engine.spp.spp = 16\n",
    "    elif engine.antialiasing_mode == 'Quasi-Random (Sobol)':\n",
    "        engine.spp.mode = 'low_discrepancy_seq'\n",
    "        engine.spp.spp = spp_slider.value\n",
    "    else:\n",
    "        raise ValueError('unknown antialiasing mode')\n",
    "        \n",
    "    engine.spp.reset_accumulation()\n",
    "    engine.use_optix_denoiser = denoiser_checkbox.value\n",
    "\n",
    "    spp_slider.value = engine.spp.spp\n",
    "    spp_slider.disabled = (engine.spp.mode == 'msaa')\n",
    "    visualizer.render_update()\n",
    "\n",
    "aa_checkbox.observe(on_change, names='value')\n",
    "aa_mode_combo.observe(on_change, names='value')\n",
    "denoiser_checkbox.observe(on_change, names='value')\n",
    "spp_slider.observe(on_change, names='value')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec1ef1dc-d7af-49c5-8594-f2ab2d0c9617",
   "metadata": {},
   "source": [
    "#### Plug 3DGRUT Engine to Canvas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "765cecac3ac0b94a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def render(in_cam, **kwargs):\n",
    "    # Called when the user stops interacting, to generate a high quality frame\n",
    "    is_first_pass = engine.is_dirty(in_cam)\n",
    "    framebuffer = engine.render_pass(in_cam, is_first_pass=True)\n",
    "    while engine.has_progressive_effects_to_render():\n",
    "        framebuffer = engine.render_pass(in_cam, is_first_pass=False)\n",
    "\n",
    "    rgba_buffer = torch.cat([framebuffer['rgb'], framebuffer['opacity']], dim=-1)\n",
    "    rgba_buffer = torch.clamp(rgba_buffer, 0.0, 1.0)\n",
    "\n",
    "    return (rgba_buffer[0] * 255).to(torch.uint8)\n",
    "\n",
    "def fast_render(in_cam, **kwargs):\n",
    "    # Called during interactions, disables effects for quick rendering\n",
    "    framebuffer = engine.render_pass(in_cam, is_first_pass=True)\n",
    "    rgba_buffer = torch.cat([framebuffer['rgb'], framebuffer['opacity']], dim=-1)\n",
    "    rgba_buffer = torch.clamp(rgba_buffer, 0.0, 1.0)\n",
    "    return (rgba_buffer[0] * 255).to(torch.uint8)\n",
    "\n",
    "# Create initial camera\n",
    "camera = kaolin.render.easy_render.default_camera(512)\n",
    "camera.change_coordinate_system(\n",
    "    torch.tensor([[1, 0, 0],\n",
    "                  [0, 0, 1],\n",
    "                  [0, -1, 0]]\n",
    "))\n",
    "camera = camera.cuda()\n",
    "# Initialize renderer\n",
    "visualizer = kaolin.visualize.IpyTurntableVisualizer(\n",
    "    height=camera.height,\n",
    "    width=camera.width,\n",
    "    camera=copy.deepcopy(camera),\n",
    "    render=render,\n",
    "    fast_render=fast_render,\n",
    "    max_fps=8,\n",
    "    world_up_axis=1\n",
    ")\n",
    "\n",
    "# Show the canvas and callback listener\n",
    "vbox = widgets.VBox([denoiser_checkbox, aa_checkbox, aa_mode_combo, spp_slider])\n",
    "hbox = widgets.HBox([visualizer.canvas, vbox])\n",
    "display(hbox, visualizer.out)\n",
    "\n",
    "# OR without a GUI, its as simple as:\n",
    "# visualizer.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3dgrt_release",
   "language": "python",
   "name": "3dgrt_release"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
